/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "npu_unwrap.h"

#include "tensorflow/core/common_runtime/eager/tensor_handle.h"

#include "npu_micros.h"
#include "npu_tensor.h"

namespace npu {
tensorflow::Status GetTensorHandleShape(TFE_TensorHandle *handle, tensorflow::TensorShape &shape) {
  tensorflow::PartialTensorShape partial_shape;
  TF_RETURN_IF_ERROR(tensorflow::unwrap(handle)->Shape(&partial_shape));
  NPU_REQUIRES(partial_shape.AsTensorShape(&shape),
               tensorflow::errors::InvalidArgument("Shape ", partial_shape.DebugString(), " not fully defined"));
  return tensorflow::Status::OK();
}

tensorflow::Status GetTensorHandleTensor(TFE_TensorHandle *handle, const tensorflow::Tensor **tensor) {
  if (IsNpuTensorHandle(handle)) {
    void *dev_buf = dynamic_cast<tensorflow::CustomDeviceTensorHandle *>(tensorflow::unwrap(handle))->DevicePointer();
    return tensorflow::TensorHandleFromInterface(
             tensorflow::unwrap(reinterpret_cast<npu::NpuTensor *>(dev_buf)->handle))
      ->Tensor(tensor);
  } else {
    return tensorflow::TensorHandleFromInterface(tensorflow::unwrap(handle))->Tensor(tensor);
  }
}
}  // namespace npu